package com.sawyerqiu.agent;

import jdk.internal.org.objectweb.asm.ClassReader;
import org.apache.commons.io.monitor.FileAlterationListenerAdaptor;
import org.apache.commons.io.monitor.FileAlterationMonitor;
import org.apache.commons.io.monitor.FileAlterationObserver;

import java.io.File;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.nio.file.Files;
import java.util.*;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

/**
 * @author : SawyerQiu
 * @date : 2024/8/1 11:40
 */
public class ClassHotLoadWatch {

    private final Instrumentation instrumentation;

    public ClassHotLoadWatch(Instrumentation instrumentation) {
        this.instrumentation = instrumentation;
    }

    public static boolean start(String libPathName, String clazzPath) throws Exception {
        Instrumentation instrumentation = new InstrumentUtil(libPathName).getInstrumentation();
        if (instrumentation == null) {
            return false;
        }
        new ClassHotLoadWatch(instrumentation).startWatch(clazzPath);
        return true;
    }

    private void startWatch(String clazzPath) throws Exception {
        File file = new File(clazzPath);
        if (!file.exists()) {
            boolean mkdir = file.mkdir();
            if (!mkdir) {
                throw new RuntimeException("创建文件夹失败：" + clazzPath);
            }
        }

        FileAlterationMonitor monitor = new FileAlterationMonitor(0);

        FileAlterationObserver observer = new FileAlterationObserver(new File(clazzPath));
        monitor.addObserver(observer);

        observer.addListener(new MyAgentFileAlterationListener());
        monitor.start();
        System.out.println("start watch:" + clazzPath);
    }

    private void refineClasses(Collection<File> files) {
        try {
            HashMap<String, byte[]> classContentMap = new HashMap<>();
            for (File file : files) {
                if (!file.getName().endsWith(".class")) {
                    continue;
                }

                byte[] bytes = Files.readAllBytes(file.toPath());
                ClassReader classReader = new ClassReader(bytes);
                String className = classReader.getClassName().replace('/', '.');
                classContentMap.put(className, bytes);
            }

            if (classContentMap.isEmpty()) {
                System.out.println("未找到需要热更的类");
                return;
            }

            ClassDefinition[] classDefinitions = new ClassDefinition[classContentMap.size()];

            int index = 0;
            for (Map.Entry<String, byte[]> entry : classContentMap.entrySet()) {
                classDefinitions[index++] = newClassDefinition(entry.getKey(), entry.getValue());
            }
            // 热更
            instrumentation.redefineClasses(classDefinitions);

            for (ClassDefinition classDefinition : classDefinitions) {
                System.out.println("热更class成功：" + classDefinition.getDefinitionClass().getName());
            }

            System.out.println("热更类完成");

        } catch (Exception e) {
            e.printStackTrace();
            System.out.println("热更类失败");
        }
    }

    private ClassDefinition newClassDefinition(String clazzName, byte[] bytes) throws Exception {
        Class<?> clazz;
        try {
            clazz = Class.forName(clazzName);
        } catch (ClassNotFoundException e) {
            clazz = defineNewClass(clazzName, bytes);
        }
        return new ClassDefinition(clazz, bytes);
    }

    private Class<?> defineNewClass(String clazzName, byte[] bytes) throws Exception {
        ClassLoader appClassLoader = Thread.currentThread().getContextClassLoader();
        System.out.println("开始新增类：" + clazzName);
        System.out.println("使用类加载器：" + appClassLoader);

        try {
            Method defineClassMethod = ClassLoader.class.getDeclaredMethod("defineClass", String.class, byte[].class, int.class, int.class);
            defineClassMethod.setAccessible(true);
            defineClassMethod.invoke(appClassLoader, clazzName, bytes, 0, bytes.length);

            Class<?> clazz = Class.forName(clazzName);
            System.out.println("完成新增类：" + clazzName);
            return clazz;
        } catch (Exception e) {
            System.out.println("新增类失败：" + clazzName);
            throw e;
        }
    }

    class MyAgentFileAlterationListener extends FileAlterationListenerAdaptor {

        private ScheduledExecutorService scheduledExecutorService;

        private Set<File> fileChangeSet;

        private ScheduledFuture scheduledFuture;

        MyAgentFileAlterationListener() {
            this.scheduledExecutorService = Executors.newSingleThreadScheduledExecutor();
            this.fileChangeSet = new HashSet<>();
        }

        @Override
        public void onFileCreate(File file) {
            execute(() -> {
                fileChangeSet.add(file);
                tryScheduleRedefineClasses();
            });
        }

        @Override
        public void onFileChange(File file) {
            execute(() -> {
                fileChangeSet.add(file);
                tryScheduleRedefineClasses();
            });
        }

        private void execute(Runnable runnable) {
            scheduledExecutorService.execute(runnable);
        }

        private void tryScheduleRedefineClasses() {
            if (scheduledFuture == null) {
                scheduledFuture = scheduledExecutorService.schedule(this::refineChangeClasses, 2000L, TimeUnit.MILLISECONDS);
            }
        }

        private void refineChangeClasses() {
            refineClasses(fileChangeSet);
            fileChangeSet.clear();

            scheduledFuture = null;
        }
    }
}
